{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1179bbd4",
   "metadata": {},
   "source": [
    "This file is for evaluating the coverage rate of generated Q&As to paper. We hope 10 answers should roughly cover the majority of the content of the paper).\n",
    "\n",
    "We calculate similarity score between each answer vs. each paper sentence. ious)\n",
    "\n",
    "We use e5-mistral-7b-instruct (paper: https://arxiv.org/pdf/2401.00368.pdf) as embedding model. It utilizes the power of LLM to capture rich contextual information of natural language. Also, it offers a more contextual-based similarity measurement through task-defining by using pro.\n",
    "\n",
    "```jsx\n",
    "task = 'Given the Question & Answer pair, retrieve relevant sentences to this answer.'\n",
    "queries = [get_detailed_instruct(task, f\"{i['Q']}\\n{i['A']}\") for i in data['Q&A']]\n",
    "documents = split_into_sentences(data['txt'])\n",
    "```\n",
    "\n",
    "By using ‘task’ as part of query, we strengthen that we want to measure the similarity between answer and paper sentence from the aspect of how much a paper sentence is related with this answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6627d2dc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(dict_keys(['source', 'doi', 'abstract', 'keywords', 'txt', 'title', 'num_Q&A', 'Q&A', 'raw_output']),\n",
       " 105)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# read file for evaluation, the file is a list of dict with keys shown in following\n",
    "# please enter your file path at 'file_path' variable\n",
    "import json\n",
    "def read(file_path):\n",
    "    with open(file_path, 'r') as file:\n",
    "        datas = json.load(file) \n",
    "    return datas\n",
    "\n",
    "file_path = '/fl/data/acl/acl_eval_final_vicuna16k_all.json'\n",
    "datas = read(file_path)\n",
    "datas[0].keys(), len(datas)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5ace292",
   "metadata": {},
   "source": [
    "### Calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9163c885",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/miniconda3/envs/xtuner/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "Loading checkpoint shards: 100%|█████████████████████████████████| 2/2 [00:08<00:00,  4.38s/it]\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch import Tensor\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "\n",
    "\n",
    "# get embedding value of the last layer\n",
    "def last_token_pool(last_hidden_states: Tensor,\n",
    "                 attention_mask: Tensor) -> Tensor:\n",
    "    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])\n",
    "    if left_padding:\n",
    "        return last_hidden_states[:, -1]\n",
    "    else:\n",
    "        sequence_lengths = attention_mask.sum(dim=1) - 1\n",
    "        batch_size = last_hidden_states.shape[0]\n",
    "        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]\n",
    "\n",
    "\n",
    "def get_detailed_instruct(task_description: str, query: str) -> str:\n",
    "    return f'Instruct: {task_description}\\nQuery: {query}'\n",
    "\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained('/fl/model/mixtral/e5-mistral-7b-instruct')\n",
    "model = AutoModel.from_pretrained('/fl/model/mixtral/e5-mistral-7b-instruct')\n",
    "model.eval()\n",
    "\n",
    "# max length should be the max token length of both queries and documents\n",
    "max_length = 1028"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "19fc61ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import gc\n",
    "import time\n",
    "\n",
    "# split paper to sentence\n",
    "def split_into_sentences(text):\n",
    "    # Regular expression for splitting sentences\n",
    "    sentence_endings = r'(?<!\\w\\.\\w.)(?<![A-Z][a-z]\\.)(?<=\\.|\\?)\\s'\n",
    "    sentences = re.split(sentence_endings, text)\n",
    "    \n",
    "    # Filter out any empty strings\n",
    "    sentences = [sentence.strip() for sentence in sentences if sentence.strip()]\n",
    "    return sentences\n",
    "\n",
    "\n",
    "# sort paper sentences' similarity score from high to low of each Q&A, recode both paper sentences' index and similarity score\n",
    "def sort_and_reformat_2d_list(list_2d):\n",
    "    sorted_2d_list = []\n",
    "    for sublist in list_2d:\n",
    "        # Sorting each sublist with scores in descending order and keeping track of original indices\n",
    "        sorted_sublist = sorted(\n",
    "            [(index, score) for index, score in enumerate(sublist)], \n",
    "            key=lambda x: x[1], \n",
    "            reverse=True\n",
    "        )\n",
    "        sorted_2d_list.append(sorted_sublist)\n",
    "    return sorted_2d_list\n",
    "\n",
    "\n",
    "# main function to calculate similarity score and record results into json file\n",
    "def para_div_mistral(data, index):\n",
    "    fn = data['doi'].replace('/', '-')\n",
    "\n",
    "    # Each query must come with a one-sentence instruction that describes the task\n",
    "    task = 'Given the Question & Answer pair, retrieve relevant sentences to this answer.'\n",
    "    queries = [get_detailed_instruct(task, f\"{i['Q']}\\n{i['A']}\") for i in data['Q&A']]\n",
    "    documents = split_into_sentences(data['txt'])\n",
    "    \n",
    "    m = len(queries)\n",
    "    n = len(documents)\n",
    "    \n",
    "    input_texts = queries + documents\n",
    "    \n",
    "    # Tokenize the input texts\n",
    "    batch_dict = tokenizer(input_texts, max_length=max_length - 1, return_attention_mask=False, padding=False, truncation=True)\n",
    "    # append eos_token_id to every input_ids\n",
    "    batch_dict['input_ids'] = [input_ids + [tokenizer.eos_token_id] for input_ids in batch_dict['input_ids']]\n",
    "    batch_dict = tokenizer.pad(batch_dict, padding=True, return_attention_mask=True, return_tensors='pt')\n",
    "\n",
    "    with torch.no_grad():\n",
    "        outputs = model(**batch_dict)\n",
    "        embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])\n",
    "    \n",
    "    # normalize embeddings\n",
    "    embeddings = F.normalize(embeddings, p=2, dim=1)\n",
    "    scores = (embeddings[:m] @ embeddings[m:].T) * 100\n",
    "    \n",
    "    similarity_dict = sort_and_reformat_2d_list(scores.tolist())\n",
    "    store = {idx:str(i) for idx, i in enumerate(similarity_dict)}\n",
    "    with open(f'{pth}/q_paragraph_sim/{fn}.json', 'w') as json_file:\n",
    "        json.dump(store, json_file, indent=4)\n",
    "\n",
    "    del batch_dict\n",
    "    del outputs\n",
    "    del embeddings\n",
    "    del scores\n",
    "    gc.collect()\n",
    "\n",
    "    if torch.cuda.is_available():\n",
    "        torch.cuda.empty_cache()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9b0f300c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "1\n",
      "2\n",
      "3\n",
      "4\n",
      "5\n",
      "6\n",
      "7\n",
      "8\n",
      "9\n",
      "10\n",
      "11\n",
      "12\n",
      "13\n",
      "14\n",
      "15\n",
      "16\n",
      "17\n",
      "18\n",
      "19\n",
      "20\n",
      "21\n",
      "22\n",
      "23\n",
      "24\n",
      "25\n",
      "26\n",
      "27\n",
      "28\n",
      "29\n",
      "30\n",
      "31\n",
      "32\n",
      "33\n",
      "34\n",
      "35\n",
      "36\n",
      "37\n",
      "38\n",
      "39\n",
      "40\n",
      "41\n",
      "42\n",
      "43\n",
      "44\n",
      "45\n",
      "46\n",
      "47\n",
      "48\n",
      "49\n",
      "50\n",
      "51\n",
      "52\n",
      "53\n",
      "54\n",
      "55\n",
      "56\n",
      "57\n",
      "58\n",
      "59\n",
      "60\n",
      "61\n",
      "62\n",
      "63\n",
      "64\n",
      "65\n",
      "66\n",
      "67\n",
      "68\n",
      "69\n",
      "70\n",
      "71\n",
      "72\n",
      "73\n",
      "74\n",
      "75\n",
      "76\n",
      "77\n",
      "78\n",
      "79\n",
      "80\n",
      "81\n",
      "82\n",
      "83\n",
      "84\n",
      "85\n",
      "86\n",
      "87\n",
      "88\n",
      "89\n",
      "90\n",
      "91\n",
      "92\n",
      "93\n",
      "94\n",
      "95\n",
      "96\n",
      "97\n",
      "98\n",
      "99\n",
      "100\n",
      "101\n",
      "102\n",
      "103\n",
      "104\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "# path to record results\n",
    "pth = 'paragraph diversity'\n",
    "\n",
    "Path(f'{pth}/q_paragraph_sim').mkdir(parents=True)\n",
    "\n",
    "for index, i in enumerate(datas):\n",
    "    print(index)\n",
    "    para_div_mistral(i, index)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a4b8aa7",
   "metadata": {},
   "source": [
    "### Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f7cf9fea",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import ast\n",
    "import numpy as np\n",
    "from scipy.stats import entropy\n",
    "\n",
    "colors = ['#FF7F50', '#40E0D0', '#6495ED', '#D8BFD8', '#98FB98', '#F0E68C', '#87CEFA', '#DDA0DD', '#7FFFD4', '#DEB887', '#5F9EA0', '#FFDAB9', '#7B68EE', '#F4A460', '#F08080']\n",
    "def load_data(file_path):\n",
    "    with open(file_path, 'r') as file:\n",
    "        data = json.load(file)\n",
    "    \n",
    "    data = {int(i): ast.literal_eval(data[i]) for i in data}    \n",
    "    return data\n",
    "\n",
    "\n",
    "# transform data from [qa, sentence] to [sentence, qa]\n",
    "def transform_data(file_path):\n",
    "    data = load_data(file_path)\n",
    "    # Find the number of paragraphs and Q&As\n",
    "    paragraph_count = len(data[0])\n",
    "    qa_count = len(data)\n",
    "\n",
    "    # Initialize a 2D array for storing similarity scores\n",
    "    similarity_matrix = np.zeros((paragraph_count, qa_count))\n",
    "\n",
    "    # Populate the similarity matrix\n",
    "    for qa_index, pairs in data.items():\n",
    "        for pair in pairs:\n",
    "            paragraph_index, sim_score = pair\n",
    "            similarity_matrix[int(paragraph_index), int(qa_index)] = sim_score\n",
    "\n",
    "    return similarity_matrix\n",
    "\n",
    "\n",
    "# std evaluation\n",
    "def calculate_metrics(similarity_matrix, save_pth):\n",
    "    std_per_paragraph = np.std(similarity_matrix, axis=1)\n",
    "    # entropy_per_paragraph = [entropy(similarity_matrix[i, :]) for i in range(len(std_per_paragraph))]\n",
    "\n",
    "    with open(save_pth, 'w') as file:\n",
    "        for s in std_per_paragraph:\n",
    "            file.write(str(s) + '\\n')\n",
    "        \n",
    "    return std_per_paragraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8e041b99",
   "metadata": {},
   "outputs": [],
   "source": [
    "# scatter plot: show the top 10 sentences’ similarity score for each answer\n",
    "def scatter_vis(similarity_dict, save_pth):\n",
    "    idx = []\n",
    "    for i in similarity_dict:\n",
    "        idx.append([j[0] for j in similarity_dict[i][:10]])\n",
    "\n",
    "    # plt.figure(figsize=(7, 4))\n",
    "    for i, indices in enumerate(idx):\n",
    "        y_values = [i+1] * len(indices)  # Same y-value for each list\n",
    "        plt.scatter(indices, y_values, color=colors[i], label=f'QA {i+1}')\n",
    "\n",
    "    plt.xlabel('Sentence Index')\n",
    "    plt.xlim(0, len(similarity_dict[0]))\n",
    "    plt.ylabel('QA Index')\n",
    "    plt.yticks(range(1, i+2, 1))\n",
    "    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')\n",
    "    plt.tight_layout()\n",
    "    # plt.show()\n",
    "    plt.savefig(save_pth)\n",
    "    plt.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6a414d97",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "# heatmap plot: show all sentences’ similarity score for each answer \n",
    "def heatmap_vis(similarity_matrix, save_pth):\n",
    "    plt.figure(figsize=(20, 3))\n",
    "    sns.heatmap(similarity_matrix.T, annot=False, cmap=\"coolwarm\")\n",
    "    plt.title(\"Heatmap of Paragraph-Q&A Similarity Scores\")\n",
    "    plt.xlabel(\"Sentence Index\")\n",
    "    plt.ylabel(\"Q&A Index\")\n",
    "    # plt.show()\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(save_pth)\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "61804b57",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm\n",
    "\n",
    "from pathlib import Path\n",
    "Path(f'{pth}/fig').mkdir(parents=True)\n",
    "Path(f'{pth}/std').mkdir(parents=True)\n",
    "\n",
    "# main function to generate plots & std scores and record them into files\n",
    "for i in range(len(datas)):\n",
    "    data = datas[i]\n",
    "    fn = data['doi'].replace('/', '-')\n",
    "    path = f'{pth}/q_paragraph_sim/{fn}.json'\n",
    "\n",
    "    scatter_pth = f'{pth}/fig/{fn}_scatter.png'\n",
    "    similarity_dict = load_data(path)\n",
    "    scatter_vis(similarity_dict, scatter_pth)\n",
    "\n",
    "    heatmap_pth = f'{pth}/fig/{fn}_heatmap.png'\n",
    "    similarity_matrix = transform_data(path)\n",
    "    heatmap_vis(similarity_matrix, heatmap_pth)\n",
    "\n",
    "    std_pth = f'{pth}/std/{fn}.txt'\n",
    "    std = calculate_metrics(similarity_matrix, std_pth)\n",
    "    # print(std)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33192e88",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "9c6cbd6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chunk_rate(sents, topk):\n",
    "    num_chunks = 10\n",
    "    sents_per_chunk = sents / num_chunks\n",
    "    chunks_hit = set()\n",
    "    for index in set(topk):\n",
    "        chunk_index = int(index // sents_per_chunk)\n",
    "        chunks_hit.add(chunk_index)\n",
    "\n",
    "    num_chunks_hit = len(chunks_hit)\n",
    "    return num_chunks_hit/num_chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "f04010e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "coverage_list = []\n",
    "for i in range(len(datas[:100])):\n",
    "    data = datas[i]\n",
    "    fn = data['doi'].replace('/', '-')\n",
    "    path = f'{pth}/q_paragraph_sim/{fn}.json'\n",
    "    sents = len(similarity_dict[0])  # This line is corrected to reflect the question's context\n",
    "    k = int(sents*0.15)\n",
    "\n",
    "    scores = [item for i in similarity_dict for item in similarity_dict[i]]\n",
    "    sorted_scores = sorted(scores, key=lambda x: x[1], reverse=True)\n",
    "    sent_idx = [d[0] for d in sorted_scores]\n",
    "\n",
    "    top15 = set(sent_idx[:k])\n",
    "    idx = k\n",
    "    while len(top15)<k:\n",
    "        top15.add(sent_idx[idx])\n",
    "        idx += 1\n",
    "    num_chunks_hit = chunk_rate(sents, top15)\n",
    "    \n",
    "    # intersection\n",
    "    top15 = []\n",
    "    individual_chunk_rate = []\n",
    "    similarity_dict = load_data(path)\n",
    "    for a in similarity_dict:\n",
    "        sent_idx = [c[0] for c in similarity_dict[a][:k]]\n",
    "        top15 += sent_idx\n",
    "        individual_chunk_rate.append(chunk_rate(sents, sent_idx))\n",
    "    \n",
    "    coverage_list.append({'doi': data['doi'], 'overall chunk rate': num_chunks_hit,\n",
    "                         'individual chunk rate': sum(individual_chunk_rate)/len(individual_chunk_rate), 'individual chunk rate list': individual_chunk_rate})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "f6288703",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(f'{pth}coverage_list.json', 'w') as json_file:\n",
    "    json.dump(coverage_list, json_file, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "bf8285ca",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.6799999999999998, 0.6423000000000001)"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "avg_chunk_rate = np.mean([i['overall chunk rate'] for i in coverage_list])\n",
    "avg_idv_chunk_rate = np.mean([i['individual chunk rate'] for i in coverage_list])\n",
    "avg_chunk_rate, avg_idv_chunk_rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "747b0364",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'paragraph diversity'"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d55a4c1",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mixtral",
   "language": "python",
   "name": "xtuner"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
